import os
import re
import fitz
from paddleocr import PaddleOCR
import math
from datetime import datetime

'''
批量识别pdf文件并添加文字层

功能：
1. 遍历文件夹目录里的所有pdf文件
2. 检查pdf文件首页是否有文字，如果没有文字，识别pdf文件并添加透明文字层
3. 识别pdf文件内容，并添加文字层
4. 保存pdf文件
'''


# 检查pdf文件首页是否有文字，
# 如果没有文字，识别pdf文件并添加透明文字层
def check_pdf_ocr_add_text(pdf_file: str):
    # 先打开 PDF 文件，获取第一页内容
    text = read_pdf_page0(pdf_file)

    # 如果text为None，开始ocr识别
    if not text:
        ocr_pdf_with_textlayer(pdf_file)
    pass


# 读取pdf内容
def read_pdf_page0(pdf_file: str):
    # 打开 PDF 文件
    with fitz.open(pdf_file) as pdf:
        page = pdf.load_page(0)  # pdf[0]
        text = page.get_text()

        if text:
            # 删除空行
            text = '\n'.join([line for line in text.split('\n') if line.strip() != ''])

            # 清除字符串中的空格和空白字符，但保留换行符
            text = re.sub(r'\s+', '', text)

            # 清除分行标记
            text = re.sub(r'\r\n+', '', text)

            # 去除矢量图形标识
            text = re.sub(r'\{#\{.*\}#\}', '', text)
            return text

    return None

# ocr识别pdf内容
def ocr_pdf_with_textlayer(pdf_file: str):
    start_time = datetime.now()
    ocr = PaddleOCR(use_angle_cls=True, lang="ch", page_num=0, use_gpu=False)
    result = ocr.ocr(pdf_file, cls=True)

    try:
        # 添加透明文字层
        with fitz.open(pdf_file) as pdf:
            for pg in range(0, pdf.page_count):
                try:
                    page = pdf[pg]
                    res = result[pg]

                    # 因为ppocr采用pdf页面长宽小于1000而采用放大策略，所以需要考虑缩放尺寸
                    scale = 0.5
                    pm = page.get_pixmap(alpha=False)
                    # if width or height > 2000 pixels, don't enlarge the image
                    if pm.width > 1000 or pm.height > 1000:
                        scale = 1.0

                    # 将识别的文本框画到pdf页面上
                    for box, [txt, score] in res:
                        # print(box, txt, score)
                        # 文字框的四边形坐标，因为ocr翻了倍数2，所以需要缩小到0.5倍
                        quad = fitz.Quad(box[0], box[1], box[2], box[3]) * scale
                        # 计算四边形近似的宽高
                        width = quad.ul.distance_to(quad.ur)
                        height = quad.ul.distance_to(quad.lr)

                        # 得到水平文字框矩形以及旋转角度
                        rect = None
                        d = None
                        if len(txt) > 1 > width/height:
                            rect = fitz.Rect(quad.ur, quad.ur.x + height, quad.ur.y + width)
                            d = quad.ll - quad.ur
                        else:
                            rect = fitz.Rect(quad.ul, quad.ul.x + width, quad.ul.y + height)
                            d = quad.ur - quad.ul

                        # 使用math.atan2计算夹角（弧度）
                        angle_rad = math.atan2(d.y, d.x)
                        # 将弧度转换为角度
                        angle_deg = math.degrees(angle_rad)

                        # 计算文本框的旋转中心点位，及角度
                        morph = (rect.tl, fitz.Matrix(1, 1).prerotate(-angle_deg))
                        fontsize = 50
                        rc = -1
                        while rc < 0:  # look for largest font size that makes the text fit
                            rc = page.insert_textbox(rect, txt,
                                                     fontname="china-ss",
                                                     fontsize=fontsize,
                                                     color=(0, 0, 0, 0),
                                                     fill_opacity=0.01,
                                                     lineheight=quad.height,
                                                     morph=morph,
                                                     )
                            fontsize -= 0.2  # 步进减小字体大小，直到匹配框大小
                    # page.commit()
                except Exception as e:
                    print(f"inner err: {e}")
            pdf.saveIncr()
    except Exception as e:
        print(f"outter err: {e}")

    end_time = datetime.now()
    time_diff = end_time - start_time
    print(f'ocr识别pdf文件：{pdf_file}, 耗时:{time_diff}')
    pass


if __name__ == '__main__':

    pdf_path = 'D:\\work\\文库资源\\2025年北京各区高三二模各科试题及答案汇总'  # 已经下载的pdf文件目录
    count = 0
    # 遍历文件夹目录里的所有pdf文件
    for root, dirs, files in os.walk(pdf_path):
        for filename in files:
            if filename.endswith('.pdf'):
                check_pdf_ocr_add_text(os.path.join(root, filename))
                count += 1
                print(f'已处理文件：{count}')
